# -*- coding: utf-8 -*-
"""
    PBAR_ConsolidateHighZMarkers.py

    - merges the high z markers found by PBAR_TestZspecBasicHighZFinder_SingleVer2.py

Created on Wed May 14 13:41:00 2014

@author: jkwong
"""

import glob, os, gc, objgraph, cPickle
import numpy as np
import PBAR_Zspec, PBAR_Cargo
import matplotlib.pyplot as plt
from scipy import ndimage
## image processing stuff
#from skimage.filter import canny
#from scipy import ndimage
#from skimage import morphology
#
#from skimage.filter import threshold_otsu
#from skimage.segmentation import clear_border
#from skimage.morphology import label, closing, square
#from skimage.measure import regionprops
#
##from skimage.color import label2rgb
#import PIL
#from PIL import Image
reload(PBAR_Zspec)

dataPath = r'E:\PBAR\data4\BasicScansStandardWidthPickle'
data0Path = r'E:\PBAR\data4\HighZFinder'
plotSaveDir = r'E:\PBAR\data4\HighZFinder'

num = 11

figureSize = (16, 10)

# Global parameters
MINBIN = 1
newWidth = 600
newHeight = 136
xCargoInternalBounds = np.array((38, 550))
yCargoInternalBounds = np.array((13, 106))

#cargoCountRange = np.array([0, 4e7]) # only consider parts of the zspec image that have rad below this value
cargoCountRange = np.array([0, 3e7]) # only consider parts of the zspec image that have rad below this value

ZspecCountRange = np.array([0, 275])

# make different size windows
# x and y
windowList = []
windowList.append(np.array((3, 1)))
windowList.append(np.array((5, 1)))
windowList.append(np.array((7, 1)))
windowList.append(np.array((4, 2)))
windowList.append(np.array((7, 2)))

#windowList.append(np.array((4, 4)))
#windowList.append(np.array((8, 8)))

discrimThresholdList = np.array((0, 2))
fractionPixelsList = np.array((0.8, 1.0))

tolerance = [3, 1]

# get linear discriminant coefficients, polynomial fits

datasetDescription = PBAR_Zspec.ReadCargoDataDescriptionFile(r'E:\PBAR\data4\CargoSet2.txt')

for datasetIndex, datasetDescript in enumerate(datasetDescription['scanID']):
    print(datasetIndex)

    if datasetIndex < 66:
        continue
    # Read Zspec
    filenameZspec = '%s-FDFC-All_SW.dat' %datasetDescription['scanID'][datasetIndex]
    fullFilenameZspec = os.path.join(dataPath, filenameZspec)
    
    # Read Cargo
    filenameCargo = 'PBAR-%s.cargoimageSW.dat' %datasetDescription['dataFile'][datasetIndex]
    fullFilenameCargo = os.path.join(dataPath, filenameCargo)

    filenameMarker = filenameCargo.replace('cargoimageSW.dat', 'cargomarkerSW')
    fullFilenameMarker = fullFilenameCargo.replace('cargoimageSW.dat', 'cargomarkerSW')
    
    energy, datZspecStandardWidth = PBAR_Zspec.ReadZspecBasicScanPickle(fullFilenameZspec)
    datZspecSmall = PBAR_Zspec.ZspecBasicReduceSize(datZspecStandardWidth, energy, newWidth)
    # calculate the features for reduced size zspec image
    discrimSmall = PBAR_Zspec.CalculateFeaturesZspecBasic(datZspecSmall, energy, MINBIN)    
    
    for windowIndex, window in enumerate(windowList):
        for discrimThresholdIndex, discrimThreshold in enumerate(discrimThresholdList):
            for fractionPixelsIndex, fractionPixels in enumerate(fractionPixelsList):
                
                highzMarkerCombinedList = []
                
                fullFilenamePotential = \
                    fullFilenameMarker.replace('cargomarkerSW', \
                        '%d_%d_%d_%3.2f.potential' %(window[0], window[1], discrimThreshold, fractionPixels) )
        
                print('Working on %s' %fullFilenamePotential)
                
                with open(fullFilenamePotential, 'rb') as fid:
                    print('Load %s' %fullFilenamePotential)
                    highzMarkerList = cPickle.load(fid)
                
                for highzMarkerIndex, highzMarker in enumerate(highzMarkerList):
                    mask = np.zeros((newWidth, newHeight)) == 1  # all falses
                    mask[highzMarker['xrange'][0]:highzMarker['xrange'][1], highzMarker['yrange'][0]:highzMarker['yrange'][1]] = True
                    if highzMarkerIndex == 0:
                        highzMarkerCombinedTemp ={}
                        highzMarkerCombinedTemp['markers'] = []
                        highzMarkerCombinedTemp['markers'].append(highzMarker)
                        highzMarkerCombinedTemp['markersIndex'] = []
                        highzMarkerCombinedTemp['markersIndex'].append(highzMarkerIndex)
                        highzMarkerCombinedTemp['mask'] = mask
                        highzMarkerCombinedTemp['maskCount'] = mask.astype(int)
                        
                        highzMarkerCombinedList.append(highzMarkerCombinedTemp)
                        continue
                    # determine which super mask most overlaps with with the current marker
                    overlapPixels = np.zeros(len(highzMarkerCombinedList))
                    for highzMarkerCombinedIndex, highzMarkerCombined in enumerate(highzMarkerCombinedList):
                        intersectionMask = highzMarkerCombined['mask'] & mask
                        # over lap in pixels
                        overlapPixels[highzMarkerCombinedIndex] = np.sum(intersectionMask.flatten())
#                    print(overlapPixels)
                    # check to see if overlap with anything
                    # if no overlap then make new marker
                    if all(overlapPixels == 0):
                        highzMarkerCombinedTemp = {}
                        highzMarkerCombinedTemp['markers'] = []
                        highzMarkerCombinedTemp['markers'].append(highzMarker)
                        highzMarkerCombinedTemp['markersIndex'] = []
                        highzMarkerCombinedTemp['markersIndex'].append(highzMarkerIndex)
                        highzMarkerCombinedTemp['mask'] = mask
                        highzMarkerCombinedTemp['maskCount'] = mask.astype(int)
                        highzMarkerCombinedList.append(highzMarkerCombinedTemp)
                    else:  # if over laps with one ore more, merge it with the one it over laps with most
                        # get the index of the super mask that this marker overlaps most with
                        maxOverlapIndex = np.argmax(overlapPixels)
                        # make union of mask
                        highzMarkerCombinedList[maxOverlapIndex]['mask'] = highzMarkerCombinedList[maxOverlapIndex]['mask'] | mask
                        highzMarkerCombinedTemp['maskCount'] = highzMarkerCombinedTemp['maskCount'] + mask.astype(int)
                        # append the marker
                        highzMarkerCombinedList[maxOverlapIndex]['markers'].append(highzMarker)
                        highzMarkerCombinedTemp['markersIndex'].append(highzMarkerIndex)
                
                # Done combining the markers, not clean up
                # - calculate mean spectrum
                # - calculat enumber of pixels
                # - calculate the features and discirminant
                # - calculate mean position
                yMatrix, xMatrix = np.meshgrid(np.arange(newHeight), np.arange(newWidth))
                for highzMarkerCombinedIndex, highzMarkerCombined in enumerate(highzMarkerCombinedList):
                    # calculate the mean x and mean y                    
                    highzMarkerCombinedList[highzMarkerCombinedIndex]['xMean'] = xMatrix[highzMarkerCombinedList[highzMarkerCombinedIndex]['mask']].flatten().mean()
                    highzMarkerCombinedList[highzMarkerCombinedIndex]['yMean'] = yMatrix[highzMarkerCombinedList[highzMarkerCombinedIndex]['mask']].flatten().mean()

                    highzMarkerCombinedList[highzMarkerCombinedIndex]['xMedian'] = np.median(xMatrix[highzMarkerCombinedList[highzMarkerCombinedIndex]['mask']].flatten())
                    highzMarkerCombinedList[highzMarkerCombinedIndex]['yMedian'] = np.median(yMatrix[highzMarkerCombinedList[highzMarkerCombinedIndex]['mask']].flatten())
                    
                    # calcalate the number of pixels
                    highzMarkerCombinedList[highzMarkerCombinedIndex]['numberPixels'] = sum(highzMarkerCombinedList[highzMarkerCombinedIndex]['mask'].flatten())
                    
                    # calculate the 
                    datZspecSmallInMaskInWindow = datZspecSmall[highzMarkerCombinedList[highzMarkerCombinedIndex]['mask'],:].mean(0)
                    discrimSmallInMaskInWindow = PBAR_Zspec.CalculateFeaturesZspecSingleSpectrum(datZspecSmallInMaskInWindow, energy, MINBIN)
                    # mean spectrum
                    highzMarkerCombinedList[highzMarkerCombinedIndex]['spectrumMean'] = datZspecSmallInMaskInWindow
                    # features, discriminant for average spectrum
                    highzMarkerCombinedList[highzMarkerCombinedIndex]['discrim'] = discrimSmallInMaskInWindow
                    
                # report summary
                print('%d markers merged into %d markers' %(len(highzMarkerList), len(highzMarkerCombinedList))  )
                # Save to file
                fullFilenameMarkerMerged = fullFilenamePotential.replace('potential', 'highZMerged')
                with open(fullFilenameMarkerMerged, 'wb') as fid:
                    print('Write %s' %fullFilenameMarkerMerged)
                    cPickle.dump(highzMarkerCombinedList, fid, 2)
                    
                # try merging with ndimage functions
                
                highzMarkerCombined2List = []
                
                if (len(highzMarkerCombinedList) > 0):  # determine if there's even anything to work with
                    mask = np.zeros_like(highzMarkerCombinedList[0]['mask']) == 1 # falses
                    for i in xrange(len(highzMarkerCombinedList)):
                        mask = mask | highzMarkerCombinedList[i]['mask']
    
                    # this finds the blobs (magic...)
                    im, number_of_objects = ndimage.label(mask)
                    
                    blobs = ndimage.find_objects(im)
                    yMatrix, xMatrix = np.meshgrid(np.arange(newHeight), np.arange(newWidth))
                    
                    # iterate through the objects
                    for blobNumber in xrange(1,number_of_objects+1):
                        highzMarkerCombined2Temp = {} # temporary container
                        subMask = im == blobNumber # mask of the contiguous group
    
                        highzMarkerCombined2Temp['mask'] = subMask
                        
                        # number of pixels in the mask
                        highzMarkerCombined2Temp['numberPixels'] = sum(subMask.flatten())
                        # calculate the mean x and mean y
                        highzMarkerCombined2Temp['xMean'] = xMatrix[subMask].flatten().mean()
                        highzMarkerCombined2Temp['yMean'] = yMatrix[subMask].flatten().mean()
    
                        highzMarkerCombined2Temp['xMedian'] = np.median(xMatrix[subMask].flatten())
                        highzMarkerCombined2Temp['yMedian'] = np.median(yMatrix[subMask].flatten())
                                            
                        # calculate the mean spectrum and features
                        datZspecSmallInMaskInWindow = datZspecSmall[subMask,:].mean(0)
                        discrimSmallInMaskInWindow = PBAR_Zspec.CalculateFeaturesZspecSingleSpectrum(datZspecSmallInMaskInWindow, energy, MINBIN)
                        # mean spectrum
                        highzMarkerCombined2Temp['spectrumMean'] = datZspecSmallInMaskInWindow
                        # features, discriminant for average spectrum
                        highzMarkerCombined2Temp['discrim'] = discrimSmallInMaskInWindow
                        
                        # keep track of original markers
                        # add orignal markers to it
                        # add original marker indices and merged marker indices so we know how this was built
                        highzMarkerCombined2Temp['markers'] = []
                        highzMarkerCombined2Temp['markersIndex'] = []
                        highzMarkerCombined2Temp['markersMergedIndex'] = []
    
                        for highzMarkerCombinedIndex, highzMarkerCombined in enumerate(highzMarkerCombinedList):
                            overlapMask = subMask & highzMarkerCombinedTemp['mask']
                            if ( overlapMask.flatten().sum() > 0):
                                highzMarkerCombined2Temp['markersIndex'] = highzMarkerCombined2Temp['markersIndex'] + highzMarkerCombined['markersIndex']
                                highzMarkerCombined2Temp['markers'] = highzMarkerCombined2Temp['markers'] + highzMarkerCombined['markers']
                                highzMarkerCombined2Temp['markersMergedIndex'].append(highzMarkerCombinedIndex)
                        
                        highzMarkerCombined2List.append(highzMarkerCombined2Temp)
                else:
                    print('No markers to worth with. Skipping and saving blank')

                # report summary
                print('%d merged-markers merged into %d merged-merged-markers' %(len(highzMarkerCombinedList), len(highzMarkerCombined2List))  )
                # Save to file
                fullFilenameMarkerMerged = fullFilenamePotential.replace('potential', 'highZMergedMerged')
                with open(fullFilenameMarkerMerged, 'wb') as fid:
                    print('Write %s' %fullFilenameMarkerMerged)
                    cPickle.dump(highzMarkerCombined2List, fid, 2)

#            break
#        break
#    break



# test
#
#
#plt.figure()
#mask = np.zeros_like(highzMarkerCombinedList[i]['mask']) == 1
#
#for i in xrange(19):
#    mask = mask | highzMarkerCombinedList[i]['mask']
#
#plt.imshow(mask.T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
#
#for i in xrange(len(highzMarkerCombinedList)):
#    plt.plot(highzMarkerCombinedList[i]['xMean'], highzMarkerCombinedList[i]['yMean'], 'xr')
#
#im, number_of_objects = ndimage.label(mask)
#blobs = ndimage.find_objects(im)
#
#plt.imsave('blobs.png', im)
#for i,j in enumerate(blobs):
#    plt.imsave('blob'+str(i)+'.png',im[j])



